using System;
using System.Buffers;
using System.Collections.Generic;
using System.Runtime.Intrinsics;
using System.Runtime.Intrinsics.X86;
using System.Threading.Tasks;

namespace ArchiVault.Core
{
    public static class ErasureCodec
    {
        #region 参数 & 初始化
        private static int _k = 10, _m = 5, _n = 15;
        public static void SetParameters(int k, int m) { _k = k; _m = m; _n = k + m; }
        #endregion

        #region GF(2^8)
        private static readonly byte[] GfLog = new byte[256];
        private static readonly byte[] GfExp = new byte[512];
        private const byte Poly = 0x1D;
        private static readonly byte[][] MulTab = new byte[256][];

        static ErasureCodec()
        {
            // 生成指数/对数表
            GfExp[0] = 1;
            for (int i = 1; i < 512; i++)
            {
                byte v = GfExp[i - 1];
                v = (byte)((v << 1) ^ (v >= 128 ? Poly : 0));
                GfExp[i] = v;
                if (i < 255) GfLog[v] = (byte)i;
            }
            // 生成乘法表
            for (int a = 0; a < 256; a++)
            {
                MulTab[a] = new byte[256];
                for (int b = 0; b < 256; b++)
                    MulTab[a][b] = (a == 0 || b == 0) ? (byte)0 : GfExp[GfLog[a] + GfLog[b]];
            }
        }
        #endregion

        #region 矩阵工具
        private static byte[,] BuildVandermonde()
        {
            var v = new byte[_n, _k];
            for (int i = 0; i < _n; i++)
                for (int j = 0; j < _k; j++)
                    v[i, j] = (byte)((i == 0 && j == 0) ? 1 : GfExp[(i * j) % 255]);
            return v;
        }

        private static void SwapRows(byte[,] a, int r1, int r2)
        {
            int c = a.GetLength(1);
            for (int j = 0; j < c; j++)
            { byte t = a[r1, j]; a[r1, j] = a[r2, j]; a[r2, j] = t; }
        }

        private static byte[,] Invert(byte[,] src)
        {
            int n = src.GetLength(0);
            var w = new byte[n, 2 * n];
            for (int i = 0; i < n; i++)
            {
                for (int j = 0; j < n; j++) w[i, j] = src[i, j];
                w[i, n + i] = 1;
            }
            for (int i = 0; i < n; i++)
            {
                if (w[i, i] == 0)
                    for (int k = i + 1; k < n; k++)
                        if (w[k, i] != 0) { SwapRows(w, i, k); break; }
                byte piv = w[i, i], invP = GfExp[255 - GfLog[piv]];
                for (int j = 0; j < 2 * n; j++) w[i, j] = MulTab[invP][w[i, j]];
                for (int k = 0; k < n; k++)
                    if (k != i && w[k, i] != 0)
                    {
                        byte f = w[k, i];
                        for (int j = 0; j < 2 * n; j++) w[k, j] ^= MulTab[f][w[i, j]];
                    }
            }
            var inv = new byte[n, n];
            for (int i = 0; i < n; i++)
                for (int j = 0; j < n; j++) inv[i, j] = w[i, n + j];
            return inv;
        }
        #endregion

        #region 核心乘法：并行 + AVX2
        private const int GRAIN = 512 * 1024;   // 512 KB 粒度

        private static byte[][] Multiply(byte[,] mat, IReadOnlyList<byte[]> inp, int sz)
        {
            int rows = mat.GetLength(0);
            int cols = mat.GetLength(1);
            // 内存池租结果
            byte[][] o = new byte[rows][];
            for (int i = 0; i < rows; i++)
                o[i] = ArrayPool<byte>.Shared.Rent(sz);

            // 按行并行
            Parallel.For(0, rows, r =>
            {
                for (int c = 0; c < cols; c++)
                {
                    byte coef = mat[r, c];
                    if (coef == 0) continue;
                    // 块处理
                    for (int offset = 0; offset < sz; offset += GRAIN)
                    {
                        int len = Math.Min(GRAIN, sz - offset);
                        unsafe
                        {
                            fixed (byte* pOut = &o[r][offset], pIn = &inp[c][offset])
                                XorBlockAvx(pOut, pIn, coef, len);
                        }
                    }
                }
            });
            return o;
        }

        private static unsafe void XorBlockAvx(byte* dst, byte* src, byte coef, int count)
        {
            if (!Avx2.IsSupported) { Fallback(); return; }

            byte* end = dst + count;
            byte* simdEnd = dst + (count & ~31);
            fixed (byte* tab = &MulTab[coef][0])
            {
                Vector256<byte> lowMask = Vector256.Create((byte)0x0F);
                while (dst < simdEnd)
                {
                    Vector256<byte> v = Avx.LoadVector256(src);
                    ReadOnlySpan<byte> tabSpan = new ReadOnlySpan<byte>(tab, 256);
                    Vector256<byte> lo = Avx2.Shuffle(Vector256.Create(tabSpan), v & lowMask);
                    Vector256<byte> hi = Avx2.Shuffle(Vector256.Create(tabSpan), Avx2.ShiftRightLogical(v.AsInt32(), 4).AsByte() & lowMask); Vector256<byte> d = Avx.LoadVector256(dst);
                    Avx.Store(dst, Avx2.Xor(d, Avx2.Xor(lo, hi)));
                    dst += 32;
                    src += 32;
                }
            }
            while (dst < end)
                *dst++ ^= MulTab[coef][*src++];

            void Fallback()
            {
                for (int i = 0; i < count; i++)
                    dst[i] ^= MulTab[coef][src[i]];
            }
        }
        #endregion

        #region 公开 API
        public static byte[][] Encode(byte[] data)
        {
            int block = (data.Length + _k - 1) / _k;
            int pad = block * _k - data.Length;
            if (pad > 0) Array.Resize(ref data, data.Length + pad);

            var slices = new byte[_k][];
            for (int i = 0; i < _k; i++)
            {
                slices[i] = ArrayPool<byte>.Shared.Rent(block);
                Array.Copy(data, i * block, slices[i], 0, block);
            }

            var result = Multiply(BuildVandermonde(), slices, block);

            // 归还数据片内存池（结果片不归还，由调用方用完再还）
            for (int i = 0; i < _k; i++)
                ArrayPool<byte>.Shared.Return(slices[i], clearArray: false);
            return result;
        }

        public static byte[] Decode(byte[][] shards)
        {
            int block = shards[0].Length;
            var sub = new byte[_k, _k];
            var vm = BuildVandermonde();
            for (int i = 0; i < _k; i++)
                for (int j = 0; j < _k; j++)
                    sub[i, j] = vm[i, j];

            var recovered = Multiply(Invert(sub), shards, block);

            byte[] data = new byte[_k * block];
            for (int i = 0; i < _k; i++)
            {
                Array.Copy(recovered[i], 0, data, i * block, block);
                ArrayPool<byte>.Shared.Return(recovered[i], clearArray: false);
            }
            return data;
        }
        #endregion
    }
}